#
# Copyright (C) [2024] Xingyun Integrated Circuit, Inc.
#
# GreenCode was a private technology asset of Xingyun Integrated Circuit， Inc （Confidential）
# Author: Shawn.Tan
# Date: 2025.10.28
#
# History: Initial Version 2025.10.28
#
#
import RandomUtils

from DV.gpgpu.trees.instruction_tree import ALU_Int64_map
from DV.gpgpu.trees.instruction_tree import LDST_Int_map
from DV.gpgpu.trees.instruction_tree import RV_G_map
from base.Sequence import Sequence
from gpgpu.EnvGPGPU import EnvGPGPU
from gpgpu.EnvGPGPU import GlobalInitSeqGPGPU
from gpgpu.GenThreadGPGPU import GenThreadGPGPU


# This class partitions the threads such that each group consists of the
# threads belonging to a given core.
class ThreadPartitionGlobalInitSeq(GlobalInitSeqGPGPU):
    def setupThreadGroup(self):
        self.partitionThreadGroup("SameCore")


# This test assigns different tasks to different groups of threads.
class MainSequence(Sequence):
    def generate(self, **kargs):
        if self.getThreadGroupId() == 0:
            self._gen_group_0_instructions()
        elif self.getThreadGroupId() == 1:
            self._gen_group_1_instructions()
        elif self.getThreadGroupId() == 2:
            self._gen_group_2_instructions()
        else:
            self._gen_shared_memory_instructions()

    # Thread Group 0 generates random integer data processing instructions.
    def _gen_group_0_instructions(self):
        for _ in range(RandomUtils.random32(50, 100)):
            instr = ALU_Int64_map.pick(self.genThread)
            self.genInstruction(instr)

    # Thread Group 1 generates random integer load/store instructions.
    def _gen_group_1_instructions(self):
        for _ in range(RandomUtils.random32(50, 100)):
            instr = LDST_Int_map.pick(self.genThread)

            # The NoPreamble flag avoids generating instructions to load a
            # value into the base register prior to generating the load or
            # store instruction
            self.genInstruction(instr, {"NoPreamble": 1})

    # Thread Group 2 randomly selects from a large collection of
    # instructions.
    def _gen_group_2_instructions(self):
        for _ in range(RandomUtils.random32(50, 100)):
            instr = RV_G_map.pick(self.genThread)
            self.genInstruction(instr)

    # Other thread groups generate load/store instructions
    # targeting a shared memory location. A thread locking context
    # permits only one thread at a time to execute until the executing
    # thread exits the context. This ensures only one thread generates
    # the shared physical address and the remaining threads use it.
    def _gen_shared_memory_instructions(self):
        shared_phys_addr = 0
        shared_phys_addr_name = "Shared PA"
        with self.threadLockingContext():
            if not self.hasSharedThreadObject(shared_phys_addr_name):
                shared_phys_addr = self.genPA(Size=8, Align=8, Type="D", Shared=1)
                self.setSharedThreadObject(shared_phys_addr_name, shared_phys_addr)
            else:
                shared_phys_addr = self.getSharedThreadObject(shared_phys_addr_name)

        target_addr = self.genVAforPA(Size=8, Align=8, Type="D", PA=shared_phys_addr)
        for _ in range(RandomUtils.random32(10, 20)):
            instr = LDST_Int_map.pick(self.genThread)
            self.genInstruction(instr, {"LSTarget": target_addr})


GlobalInitialSequenceClass = ThreadPartitionGlobalInitSeq
MainSequenceClass = MainSequence
GenThreadClass = GenThreadGPGPU
EnvClass = EnvGPGPU
